import boto3
import time

# 配置参数（请根据实际情况修改）
ATHENA_DATABASE = 'your_database_name'  # Athena 数据库名
ATHENA_OUTPUT = 's3://your-output-bucket/athena/results/'  # 查询结果输出到 S3 的路径
REGION = 'us-east-1'  # AWS 区域


def run_athena_query(query, database, output, region):
    client = boto3.client('athena', region_name=region)
    response = client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={'OutputLocation': output}
    )
    return response['QueryExecutionId']


def wait_for_query(query_execution_id, client):
    while True:
        response = client.get_query_execution(QueryExecutionId=query_execution_id)
        state = response['QueryExecution']['Status']['State']
        if state in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            return state
        time.sleep(2)


def get_query_results(query_execution_id, client):
    response = client.get_query_results(QueryExecutionId=query_execution_id)
    return response['ResultSet']


if __name__ == '__main__':
    # 你要执行的SQL（请替换为你的表名）
    sql = 'SELECT * FROM your_table LIMIT 10;'
    client = boto3.client('athena', region_name=REGION)
    print("提交查询...")
    query_id = run_athena_query(sql, ATHENA_DATABASE, ATHENA_OUTPUT, REGION)
    print(f"查询ID: {query_id}")
    print("等待查询完成...")
    state = wait_for_query(query_id, client)
    print(f"查询状态: {state}")
    if state == 'SUCCEEDED':
        results = get_query_results(query_id, client)
        print("查询结果：")
        for row in results['Rows']:
            print([col.get('VarCharValue', '') for col in row['Data']])
    else:
        print("查询未成功，请检查SQL或Athena配置。")
